#include "xict_common/utils/exposure_compensation.h"
#include <numeric>

namespace xict_common {
    namespace utils {
        std::vector<float> ExposureCompensation::CalculateHistogramThreshold(
            const std::vector<cv::Mat>& na, int pixel_num) {
            std::vector<float> threshold(3);

            // Calculate the threshold for each channel
            for (int c = 0; c < 3; ++c) {
                // Create a sorted list of histogram values with their indices
                std::vector<std::pair<float, int>> hist_sorted(256);
                for (int i = 0; i < 256; ++i) {
                    hist_sorted[i] = std::make_pair(na[c].at<float>(i), i);
                }
                std::sort(hist_sorted.begin(), hist_sorted.end());

                // Calculate the cumulative sum of histogram values
                float sum  = 0;
                int number = 0;
                // Find the threshold that corresponds to 5% of the total pixel
                // count
                while (sum < pixel_num / 20.0f && number < 256) {
                    sum += hist_sorted[number].first;
                    number++;
                }

                // Set the threshold to the histogram value at the found index
                threshold[c] = na[c].at<float>(hist_sorted[number - 1].second);
            }

            return threshold;
        }

        cv::Mat ExposureCompensation::ApplyHistogramMapping(
            const std::vector<cv::Mat>& map, const cv::Mat& input) {
            cv::Mat output = input.clone();
            int h          = input.rows;
            int w          = input.cols;

            // Apply the mapping to each pixel in each channel
            for (int c = 0; c < 3; ++c) {
                for (int i = 0; i < h; ++i) {
                    for (int j = 0; j < w; ++j) {
                        // Get the current pixel value
                        int val = input.at<cv::Vec3b>(i, j)[c];
                        // Apply the mapping and store the result
                        output.at<cv::Vec3b>(i, j)[c] =
                            static_cast<uchar>(map[c].at<float>(val) - 1);
                    }
                }
            }

            return output;
        }

        std::vector<cv::Mat>
        ExposureCompensation::PerformHybridHistogramMatching(
            const cv::Mat& im_tgt, const cv::Mat& im_src) {
            // 0. Prepare
            int H         = im_src.rows;
            int W         = im_src.cols;
            int pixel_num = H * W;

            // Resize target image to match source image dimensions
            cv::Mat im_tgt_resized;
            cv::resize(im_tgt, im_tgt_resized, cv::Size(W, H));

            // 1. Calculates histograms and cumulative histograms for both
            // images Initialize vectors to store histograms and cumulative
            // histograms
            std::vector<cv::Mat> na(3), nb(3), Sa(3), Sb(3);
            // Calculate histograms and cumulative histograms for each channel
            for (int c = 0; c < 3; ++c) {
                std::vector<cv::Mat> channels_tgt, channels_src;
                cv::split(im_tgt_resized, channels_tgt);
                cv::split(im_src, channels_src);

                // Calculate histograms
                cv::calcHist(std::vector<cv::Mat> {channels_tgt[c]}, {0},
                             cv::Mat(), na[c], {256}, {0, 256});
                cv::calcHist(std::vector<cv::Mat> {channels_src[c]}, {0},
                             cv::Mat(), nb[c], {256}, {0, 256});

                // Calculate cumulative histograms
                Sa[c] = na[c].clone();
                Sb[c] = nb[c].clone();
                for (int n = 1; n < 256; ++n) {
                    Sa[c].at<float>(n) += Sa[c].at<float>(n - 1);
                    Sb[c].at<float>(n) += Sb[c].at<float>(n - 1);
                }
            }

            // 2. Computes an initial mapping based on cumulative histograms
            // Initialize mapping and index matrices
            std::vector<cv::Mat> map(3);
            std::vector<cv::Mat> index(3);
            for (int c = 0; c < 3; ++c) {
                map[c]   = cv::Mat::zeros(1, 256, CV_32F);
                index[c] = cv::Mat::ones(1, 256, CV_8U);
            }

            // 3. Applies thresholding to remove noise
            // Calculate thresholds for histogram pruning
            std::vector<float> threshold_a =
                CalculateHistogramThreshold(na, pixel_num);
            std::vector<float> threshold_b =
                CalculateHistogramThreshold(nb, pixel_num);

            // Main mapping calculation loop
            for (int c = 0; c < 3; ++c) {
                // Find global min and max values across all channels
                double srcMin, srcMax;
                for (int c = 0; c < 3; ++c) {
                    cv::Mat channel;
                    cv::extractChannel(im_src, channel, c);
                    cv::minMaxLoc(channel, &srcMin, &srcMax);
                    srcMax = std::max(srcMax, srcMax);
                    srcMin = std::min(srcMin, srcMin);
                }
                srcMax += 1;
                srcMin += 1;

                // Calculate mapping for each intensity level
                for (int a = 0; a < 256; ++a) {
                    int b = 0;
                    while (Sa[c].at<float>(a) > Sb[c].at<float>(b)) {
                        b++;
                        if (b >= srcMax) {
                            b = srcMax;
                            break;
                        }
                    }
                    if (b < srcMin) {
                        b = srcMin;
                    }
                    map[c].at<float>(a) = b;

                    // Apply thresholding to remove noise
                    if (na[c].at<float>(a) < threshold_a[c]) {
                        map[c].at<float>(a) =
                            std::numeric_limits<float>::quiet_NaN();
                        index[c].at<uchar>(a) = 0;
                    }
                    if (nb[c].at<float>(b) < threshold_b[c]) {
                        map[c].at<float>(a) =
                            std::numeric_limits<float>::quiet_NaN();
                        index[c].at<uchar>(a) = 0;
                    }
                }

                // Set boundary conditions
                map[c].at<float>(0)     = srcMin;
                map[c].at<float>(255)   = srcMax;
                index[c].at<uchar>(0)   = 1;
                index[c].at<uchar>(255) = 1;
            }

            // 4. Performs local adjustments to preserve image details
            // Gradient and region filling logic
            for (int c = 0; c < 3; ++c) {
                // Calculate gradient
                std::vector<float> gradient(256, 0);
                for (int a = 1; a < 256; ++a) {
                    gradient[a] =
                        index[c].at<uchar>(a) - index[c].at<uchar>(a - 1);
                }

                // Identify regions to be filled
                std::vector<std::vector<int>> region;
                int n_re = 0;
                for (int a = 0; a < 256; ++a) {
                    if (gradient[a] == -1) {
                        region.push_back({a, 0});
                        n_re++;
                    } else if (gradient[a] == 1 && n_re > 0) {
                        region[n_re - 1][1] = a - 1;
                    }
                }

                // Fill identified regions using linear interpolation
                for (const auto& r : region) {
                    int start = r[0];
                    int end   = r[1];

                    std::vector<float> XX = {static_cast<float>(start - 1),
                                             static_cast<float>(end + 1)};
                    std::vector<float> YY = {map[c].at<float>(start - 1),
                                             map[c].at<float>(end + 1)};

                    auto p = polyfit(XX, YY);

                    for (int a = start; a <= end; ++a) {
                        map[c].at<float>(a) = polyval(p, a);
                    }
                }
            }

            // 5. Smooths the final mapping
            // Smooth the mapping using Gaussian blur
            for (int c = 0; c < 3; ++c) {
                cv::GaussianBlur(map[c], map[c], cv::Size(15, 1), 0);
            }

            return map;
        }

        bool ExposureCompensation::CompensateExposure(
            const cv::Mat& img1, const cv::Mat& img2, cv::Mat& img1_compensated,
            cv::Mat& img2_compensated, const cv::Point2i& offset) {
            // Calculate histogram matching map and compensate
            std::vector<cv::Mat> map =
                PerformHybridHistogramMatching(img2, img1);
            cv::Mat correct = ApplyHistogramMapping(map, img2);

            // Output compensated images
            img1_compensated = img1;
            img2_compensated = correct;

            return true;
        }
    }   // namespace utils
}   // namespace xict_common